rocm: fix native extensions, kernel JIT builds, and Triton PTX fallbacks for gfx1150 - #241
rocm: fix native extensions, kernel JIT builds, and Triton PTX fallbacks for gfx1150#241skywalk1411 wants to merge 4 commits into
Conversation
…cks for gfx1150
FreeToken built and ran only against CUDA. On a native ROCm install (tested
on a Ryzen AI 9 HX 470 / Radeon 890M, gfx1150) it failed at every stage:
build, JIT compile, and finally a silent native crash mid-warmup with no
Python traceback.
Build system (setup.py, hip_compat.h):
- _pinned_tensor and _cpu_moe link against HIP instead of cudart when
ROCM_HOME is present (CUDA_HOME stays required on the CUDA path).
- kernel/csrc/hip_compat.h aliases the CUDA Runtime API calls those two
files use onto their HIP equivalents, including CUDART_CB (undefined
under HIP, which otherwise corrupts the surrounding declaration's parse).
- The pip-vendored ROCm SDK ships versioned sonames (libamdhip64.so.7) with
no bare .so dev symlink, so link the exact file via -l:; the dynamic
linker dedupes by SONAME at runtime against whatever libamdhip64 torch
itself already loaded.
Shared kernel header (kernel/csrc/include/freetoken/utils.cuh):
- Explicit HIP runtime include + CUDA Runtime API aliases (nvcc pulls
cuda_runtime.h in implicitly for .cu files; hipcc does not).
- __grid_constant__ has no HIP equivalent; falls back to an ordinary
by-value kernel parameter.
- LaunchKernel has no HIP path for cudaLaunchKernelEx/cudaLaunchConfig_t
(that API only exists to carry Hopper PDL attributes) -- added a HIP
variant that launches via plain triple-chevron syntax instead, with
with_attr() as a no-op since there is no attribute to carry.
- The griddepcontrol PDL asm is now unconditionally a no-op under HIP,
not just when kUsePDL is false, so a stray HIP-side call can't try to
assemble Hopper-only PTX.
Kernel JIT (kernel/utils.py):
- Drop --expt-relaxed-constexpr on HIP; hipcc/clang rejects it outright.
Triton kernels:
- norm.py, activation.py: launch_pdl is a CUDA-Hopper-only kwarg; the AMD
arg-packer raises KeyError on it even when passed as False, so it's only
included when pdl is actually true (never on ROCm).
- attention.py: the decode kernel's GQA head-tile floors at 16 under HIP
(RDNA WMMA has no instruction below M=16); the kernel already masks
padded head lanes for non-power-of-two groups, so this is a safe
widening. Falls back to broadcast-multiply-reduce instead of tl.dot for
that tile as a second-layer guard. The split extend/prefill kernel's
tile shrinks from 128x64 to 64x32 under HIP -- running both the cached-
and newly-computed-KV loops live at once is register-heavier than the
plain extend kernel, and exhausts this GPU's VGPR file at the CUDA-tuned
tile size.
- activation.py (the actual root cause of the crash above): _fast_tanh and
_fast_ex2 inline raw PTX text (tanh.approx.f32, ex2.approx.f32) via
tl.inline_asm_elementwise. HIP's inline-asm path doesn't reject foreign
PTX at parse time -- it fails much later in register allocation with a
generic, misleading diagnostic ("couldn't allocate output register for
constraint 'f'") that looks like a matrix-core or register-pressure
issue and sent debugging down that path for a while. Routed through
libdevice.tanh / tl.exp2 on HIP instead.
pyproject.toml: loosen the torch/triton ceilings so ROCm builds (which
carry a local version segment such as +rocm7.14.0...) can satisfy them.
Every change is gated on HIP detection (torch.version.hip / ROCM_HOME /
__HIP_PLATFORM_AMD__) at build or run time; the CUDA path is unchanged.
Verified end to end on gfx1150: server boot, weight load, KV cache alloc,
CUDA graph capture at bs=1/2/4, and real chat completions against
Qwen/Qwen3-8B (bf16, triton attention backend).
…t-copy kernel Two more real bugs found while running an actual MoE model (Qwen3.6-35B-A3B-FP8, --moe-backend offload) end to end on gfx1150, past what the first commit covered. e4m3_compat.py: e4m3_native() decides whether kernels get raw fp8 tensors or a uint8 view by checking torch.cuda.get_device_capability() >= (8, 9). On a HIP build that call returns the GPU's RDNA generation number, not a CUDA compute capability -- gfx1150 reports (11, 5), and (11, 5) >= (8, 9) is True by plain tuple comparison (11 > 8), so this incorrectly claimed native fp8 support on AMD. Triton's own compile-time twin, e4m3_native_cx() (target_info. cuda_capability_geq, which checks target.backend != "cuda" first), correctly said False, so the kernel compiled for the emulated uint8 path while the host side hands it an untouched fp8 tensor -- IncompatibleTypeErrorImpl inside e4m3_u8_to_f32's bitwise ops. Fixed by checking torch.version.hip first. fast_index_copy.cuh (the offload cache's fast host->device expert-copy kernel, only exercised once a real MoE model with --moe-backend offload actually streams experts): same two problems as the first commit's fixes elsewhere in this file family, just not caught until this path actually ran. - Missing HIP aliases for cudaGetDevice/cudaDeviceGetAttribute/ cudaHostGetDevicePointer/the two cudaDevAttr* constants it uses -- added to utils.cuh's existing HIP block alongside the ones from the first commit. - load_nc/store_nc inline raw PTX (ld.global.L1::no_allocate, st.global.wt -- cache-policy hints, no HIP equivalent). Falls back to plain loads/stores under HIP; correctness unchanged, only the cache hint is lost. Verified: Qwen3.6-35B-A3B-FP8 (256 experts/layer x 40 layers, 3B active) boots and serves real chat completions with --moe-backend offload --moe-cache-size 2560 (25% of the model's 10240 total experts resident, LRU-evicting the rest from host RAM on every miss) -- ft ctl cache confirms the pool is live at the requested size, not silently falling back to full residency.
|
Pushed a second commit — two more real bugs, found by actually running a MoE model end to end (the first commit only covered the dense Qwen3-8B boot path):
Verified against |
Third loading path verified: google/gemma-4-26B-A4B-it-qat-q4_0-gguf (native GGUF, MoE offload) now boots and serves on gfx1150, alongside the dense bf16 and FP8 MoE paths from the earlier commits. kernel/gguf.py: same nvcc-only-flag problem as elsewhere in this port, in a third JIT mechanism (torch.utils.cpp_extension.load, distinct from both setup.py's CppExtension and the tvm-ffi JIT the rest of kernel/ uses). --expt-relaxed-constexpr is rejected outright, and the -ccbin/CXX-forcing block exists only to work around an nvcc+libtorch-headers compiler mismatch that doesn't apply under hipcc (its own bundled clang already is the host compiler). Both dropped on HIP. kernel/csrc/gguf/dispatch.h: the donor's SGLANG_SHFL_XOR_SYNC(_WIDTH) macros forward a CUDA-style 32-bit mask straight into __shfl_xor_sync. HIP's amd_warp_sync_functions.h static_asserts the mask must be 64 bits unconditionally (regardless of actual wavefront width) -- widened the cast on HIP only. .gitignore: torch's ROCm auto-hipify (a real, working translation pass built into torch.utils.cpp_extension -- unlike the other two JIT paths, this one needed no manual porting for the .cu/.cuh sources themselves) writes translated copies next to the CUDA sources it processes (gguf_kernel.cu -> .hip, *.cuh -> *_hip.cuh). Regenerated every build, never hand-edited; ignore rather than track. Not in this commit, environment-only: the pip ROCm nightly distribution used here (rocm.nightlies.amd.com) ships no thrust/rocprim headers, which torch's own extension headers pull in transitively. Ubuntu's librocthrust-dev is one fix, but it depends on libamdhip64-dev, which drops a second, conflicting HIP header set into /usr/include/hip that silently wins over the correct pip-bundled ones for any plain -I (though not -isystem) -- diagnosed by hand with `clang++ -v` and a minimal reproducer. Worked around locally by extracting just the thrust/rocprim headers (dpkg -x, no install) into the pip package's own include dir and removing the conflicting system packages; ROCM_PATH/HIP_PATH/HIP_DEVICE_LIB_PATH also had to point at the pip package for this JIT path's device-bitcode-library lookup. Left out of the diff since there's no source change to make -- noting it here for the next person on this distribution.
|
Pushed a third commit — GGUF now works too ( Two small source fixes, same pattern as before:
Worth calling out: this JIT path uses torch's own ROCm auto-hipify ( One thing that isn't in the diff, flagged for the next person on this distribution: the pip ROCm nightly build ( Test plan checklist updated. Only multi-GPU (pynccl) is left untested. |
Both were made mid-investigation, before the real cause of a since-fixed crash (the raw-PTX bug in activation.py, and separately the e4m3_native() tuple- comparison bug) was actually found. Re-tested each in isolation -- eager, batched, and inside real CUDA graph capture+replay -- now that those are fixed, and both work fine at the original, CUDA-tuned settings: - decode_paged_attention: the block_h>=16 floor (kept -- RDNA WMMA genuinely has no instruction below M=16, confirmed independently and matches upstream FlashML-org#137) was sufficient on its own. The USE_TL_DOT broadcast-sum fallback this PR had added on top was solving a problem that was actually in a different kernel; removed, restoring real matrix-core-accelerated decode attention. - _select_extend_tile: the 128x64 -> 64x32 shrink on HIP was diagnosed as a VGPR-exhaustion issue via a py-spy trace mid-investigation, before the session had isolated the actual crash to activation.py. Re-verified end-to-end against Qwen3.6-35B-A3B-FP8's GDN/split-extend path (the kernel this shrink targeted) at the original tile size: no crash, correct output. Reverted to the CUDA-tuned tile. Both re-verified against real chat completions (Qwen3-8B for the decode path, Qwen3.6-35B-A3B-FP8 for the extend/split path) after reverting, not just the isolated kernel tests.
|
Pushed a fourth commit — performance follow-up, no new hardware coverage this time. Ran That prompted a second look at two defensive fixes from the earlier commits, both made mid-investigation before the real crash cause (the
Both changes reduce the diff (net -26 lines) and restore full matrix-core acceleration + the original tile size rather than leaving unnecessary slow-path code in place. Re-verified against real chat completions on both models after reverting, not just the isolated kernel tests. |
1 participant
Summary
Gets FreeToken booting and serving on native ROCm (tested on a Ryzen AI 9 HX 470 / Radeon 890M,
gfx1150). Verified end to end: server boot, weight load, KV cache allocation, CUDA graph capture at bs=1/2/4, and real chat completions againstQwen/Qwen3-8B(bf16, triton attention backend).This covers the same ground as #137 (build system, kernel JIT flags,
launch_pdl, the decode-attention WMMA tile floor) and independently arrived at the same fixes there, plus two more that block a working boot on this GPU:kernel/triton/activation.py's_fast_tanh/_fast_ex2inline raw NVIDIA PTX text (tanh.approx.f32,ex2.approx.f32) viatl.inline_asm_elementwise. HIP's inline-asm path doesn't reject foreign PTX at parse time — it fails much later in register allocation with a generic, misleading diagnostic (error: couldn't allocate output register for constraint 'f') that looks like a matrix-core or register-pressure issue. This is called from every MLP layer's SiLU, so it reproduces on any model, deterministically, and took the longest to isolate (neededpy-spylive process inspection to catch the real call site, since the crash bypasses Python's exception handling entirely). Routed throughlibdevice.tanh/tl.exp2on HIP instead._extend_attention_split_kernel) runs both the cached-KV and newly-computed-KV loops live at once, which is register-heavier than the plain extend kernel and exhausts this GPU's VGPR file at the CUDA-tuned 128x64 tile. Shrinks to 64x32 under HIP.Every change is gated on HIP detection (
torch.version.hip/ROCM_HOME/__HIP_PLATFORM_AMD__) at build or run time; the CUDA path is unchanged.Details by file
Build system (
setup.py, newkernel/csrc/hip_compat.h)_pinned_tensor/_cpu_moelink against HIP instead of cudart whenROCM_HOMEis present.hip_compat.haliases the CUDA Runtime API calls those two files use onto HIP, includingCUDART_CB(undefined under HIP, which otherwise corrupts the surrounding declaration's parse).libamdhip64.so.7) with no bare.sodev symlink, so link the exact file via-l:; the dynamic linker dedupes by SONAME at runtime against whateverlibamdhip64torch itself already loaded.Shared kernel header (
kernel/csrc/include/freetoken/utils.cuh)cuda_runtime.himplicitly for.cufiles; hipcc does not).__grid_constant__has no HIP equivalent; falls back to an ordinary by-value kernel parameter.LaunchKernelhas no HIP path forcudaLaunchKernelEx/cudaLaunchConfig_t(that API only exists to carry Hopper PDL attributes) — added a HIP variant that launches via plain triple-chevron syntax, withwith_attr()as a no-op.griddepcontrolPDL asm is now unconditionally a no-op under HIP, not just whenkUsePDLis false, so a stray HIP-side call can't try to assemble Hopper-only PTX.Kernel JIT (
kernel/utils.py): drop--expt-relaxed-constexpron HIP; hipcc/clang rejects it outright.Triton kernels
norm.py,activation.py:launch_pdlis a CUDA-Hopper-only kwarg; the AMD arg-packer raisesKeyErroron it even when passed asFalse, so it's only included whenpdlis actually true.attention.py: decode kernel's GQA head-tile floors at 16 under HIP (RDNA WMMA has no instruction below M=16); falls back to broadcast-multiply-reduce instead oftl.dotfor that tile as a second-layer guard. Extend/prefill split kernel tile shrinks 128x64 → 64x32 under HIP (see above).activation.py:_fast_tanh/_fast_ex2PTX →libdevice.tanh/tl.exp2on HIP (see above).pyproject.toml: loosen the torch/triton version ceilings so ROCm builds (which carry a local version segment such as+rocm7.14.0...) can satisfy them.Test plan
pip install -e .builds_pinned_tensor/_cpu_moecleanly against ROCmft serve --model Qwen/Qwen3-8Bboots, loads weights, allocates KV cache/v1/chat/completionsreturns correct, coherent outputQwen/Qwen3.6-35B-A3B-FP8(256 experts/layer x 40 layers, 3B active) with--moe-backend offload --moe-cache-size 2560(25% of experts resident, LRU-evicting the rest from host RAM) -- boots, captures CUDA graphs, serves correct chat completions,ft ctl cacheconfirms the pool is live at the requested sizegoogle/gemma-4-26B-A4B-it-qat-q4_0-gguf(native GGUF dequant kernels, MoE offload) -- boots and serves correct chat completions. See the third commit for the source fixes and a documented environment gotcha (pip ROCm nightly ships no thrust/rocprim headers; the obvious apt fix pulls in a second, conflicting HIP header set)